{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quaternions representing rotations SO(3)\n",
    "This notebook is a short sumamry of how SO(3) is handled in Pinocchio, and mostly quaternion representation. \n",
    "## SO(3), Euler angles, rotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import magic_donotload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "R =\n",
      "[[-0.27142346  0.44068496 -0.85564366]\n",
      " [-0.7458126   0.46562786  0.47639718]\n",
      " [ 0.6083526   0.76745519  0.20228607]]\n",
      "p =\n",
      "[ 0.21393775 -0.96739886 -0.51422646]\n"
     ]
    }
   ],
   "source": [
    "import pinocchio\n",
    "from pinocchio import SE3, Quaternion\n",
    "import numpy as np\n",
    "from numpy.linalg import norm\n",
    "M = SE3.Random()\n",
    "R = M.rotation\n",
    "p = M.translation\n",
    "print('R =\\n' + str(R))\n",
    "print('p =\\n' + str(p.T))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A rotation is simply a 3x3 matrix. It has a unit norm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 1.00000000e+00 -1.38672842e-16  4.65901581e-17]\n",
      " [-1.38672842e-16  1.00000000e+00 -9.39205724e-17]\n",
      " [ 4.65901581e-17 -9.39205724e-17  1.00000000e+00]]\n"
     ]
    }
   ],
   "source": [
    "print(R @ R.T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It can be equivalently represented by a quaternion. Here we have made the choice to create a specific class for quaternions (i.e. they are not vectors, and can be e.g. multiplied), but you can get the 4 coefficients with the adequate method. Note that the corresponding vector is also of norm 1.\n",
    "\n",
    "它可以用一个四元数等效表示。在这里，我们选择为四元数创建一个特定的类（即，它们不是向量，例如可以进行乘法运算），但你可以通过适当的方法获得这4个系数。请注意，相应的向量范数也为1。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0\n"
     ]
    }
   ],
   "source": [
    "quat = Quaternion(R)\n",
    "print(norm(quat.coeffs()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Angle-axis representation are also implemented in the class AngleAxis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.877329005764724 [ 0.15264444 -0.76778813 -0.62225483]\n"
     ]
    }
   ],
   "source": [
    "from pinocchio import AngleAxis\n",
    "utheta = AngleAxis(quat)\n",
    "print(utheta.angle, utheta.axis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can display rotation in the viewer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You can open the visualizer by visiting the following URL:\n",
      "http://127.0.0.1:7004/static/\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "            <div style=\"height: 400px; width: 100%; overflow-x: auto; overflow-y: hidden; resize: both\">\n",
       "            <iframe src=\"http://127.0.0.1:7004/static/\" style=\"width: 100%; height: 100%; border: none\"></iframe>\n",
       "            </div>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from utils.meshcat_viewer_wrapper import MeshcatVisualizer\n",
    "viz = MeshcatVisualizer()\n",
    "viz.viewer.jupyter_cell()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz.addBox('world/box', [.1, .2, .3], [1, 0, 0, 1])\n",
    "viz.applyConfiguration('world/box', [.1, .2, .3] + list(quat.coeffs()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quaternion you said?\n",
    "Quaternions are \"complex of complex\", introduced form complex as complex are from reals. Let's try to understand what they contains in practice. Quaternions are of the form w+xi+yj+zk, with w,x,y,z real values, and i,j,k the 3 imaginary numbers. We store them as 4-d vectors, with the real part first: quat = [x,y,z,w]. We can interprete w as encoding the cosinus of the rotation angle. Let's see that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.938664502882362\n",
      "1.877329005764724\n"
     ]
    }
   ],
   "source": [
    "from numpy import arccos\n",
    "print(arccos(quat[3]))\n",
    "print(AngleAxis(quat).angle)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indeed, w = cos(θ/2). Why divided by two? For that, let's see how the quaternion can be used to represent a rotation. We can encode a 3D vector in the imaginary part of a quaternion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(x,y,z,w) = 0.562641 0.746494 0.821965        0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "p = np.random.rand(3)\n",
    "qp = Quaternion(0., p[0], p[1], p[2])\n",
    "print(qp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The real product extends over quaternions, so let's try to multiply quat with p:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(x,y,z,w) = 0.198049 0.057399 0.926117 0.805751\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(quat * qp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Well that's not a pure imaginary quaternion anymore. And the imaginary part does not contains somethig that looks like the rotated point:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.52705421  0.31954568  1.08145669]\n"
     ]
    }
   ],
   "source": [
    "print(quat.matrix() @ p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pure quaternion is obtained by multiplying again on the left by the conjugate (w,-x,-y,-z)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(x,y,z,w) = -0.527054  0.319546   1.08146         0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(quat * qp * quat.conjugate())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That is a pure quaternion, hence encoding a point, and does corresponds to R*p. Magic, is it not? We can prove that the double product of quaternion does corresponds to the rotation. Indeed, a quaternion rather encode an action (a rotation) in $R^4$, but which moves our point p outside of $R^3$. The conjugate rotation brings it back in $R^3$ but applies a second rotation. Since we rotate twice, it is necessary to apply only half of the angle each time.\n",
    "What if we try to apply the rotation quat on the imaginary part of the quaternion?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(x,y,z,w) =  0.123149 -0.619428 -0.502016         0\n",
      " (x,y,z,w) =    0.123149   -0.619428   -0.502016 2.77556e-17\n",
      "\n"
     ]
    }
   ],
   "source": [
    "qim = Quaternion(quat)  # copy\n",
    "qim[3] = 0\n",
    "print(qim, quat * qim * quat.conjugate())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What kind of conclusion can we get from this? What geometrical interpretation can we give to $q_{im}$? What about $||q_{im}||$?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The SLERP example\n",
    "Let's practice! Implement a linear interpolation between two position p0 and p1, i.e. find the position p(t) with t varying from 0 to 1, with p(0)=p0, p(1)=p1 and continuity between the two extrema.\n",
    "\n",
    "让我们练习一下！在两个位置 p0 和 p1 之间进行线性插值，也就是说，找到位置 p(t)，其中 t 的取值范围是从 0 到 1，且 p(0)=p0，p(1)=p1，并且在两个极值之间保持连续性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the solution should you need it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "%do_not_load appendix/solution_lerp.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LERP with quaternions is not working because they are not normalize. Instead we can take either the normalization of the LERP (NLERP), or the spherical LERP (SLERP). \n",
    "\n",
    "使用四元数进行线性插值（LERP）不起作用，因为它们没有进行归一化。相反，我们可以采用线性插值的归一化形式（NLERP），或者球面线性插值（SLERP）。 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the solution if you need it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "%do_not_load appendix/solution_slerp.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cheat sheet\n",
    "\n",
    "You can convert quaternion to rotation matrix and create SE3 objects as follows:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "qu = Quaternion(.7,.2,.2,.6)# Quaternion: take care that norm <= 1 (and approx 1)\n",
    "R  = qu.matrix()                # Create a rotation matrix from quaternion\n",
    "p  = np.array([0.,0.,0.77])     # Translation (R3) vector)\n",
    "M  = SE3(R,p)                   # Create a nomogeneous matrix from R,P\n",
    "\n",
    "# Typical tool position\n",
    "from pinocchio.utils import rotate\n",
    "M = SE3(rotate('z',1.)@rotate('x',.2), np.array([0.1,0.02,.65]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "robotics_course",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
